766E - Mahmoud and a xor trip - CodeForces Solution

#include <bits/stdc++.h>
// #include <ext/pb_ds/assoc_container.hpp>
// #include <ext/pb_ds/tree_po?licy.hpp>

using namespace std;
// using namespace __gnu_pbds;

#pragma GCC optimize ("unroll-loops,Ofast,O3")
#pragma GCC optimize("fast-math")
#pragma GCC optimize("no-stack-protector")
#pragma GCC target("sse,sse2,sse3,ssse3,popcnt,abm,mmx,tune=native")

#define F first
#define S second
#define sz(x) (int)x.size()
#define pb push_back
#define eb emplace_back
#define all(x) (x).begin(),(x).end()
#define rall(x) (x).rbegin(),(x).rend()
#define NFS ios_base::sync_with_stdio(0) , cin.tie(0) , cout.tie(0) ;
// #define file(s) if (fopen(s".in", "r")) freopen(s".in", "r", stdin), freopen(s".out", "w", stdout)
// #define int long long

typedef long long ll;


const int N = 1e5 + 7;
const int L = log2(1e6);
int n, k, bit[L+2][3], s[N], w[N];
bool used[N];
long long res;
vector<int> g[N];

void calc_size(int v, int p) {
    s[v] = 1;
    for (int u : g[v])
        if (u != p && !used[u])
            calc_size(u, v), s[v] += s[u];
}

int centroid(int v, int p, int n) {
    for (int u : g[v])
        if (u != p && !used[u] && s[u] > n/2)
            return centroid(u, v, n);
    return v;
}

void dfs (int v, int p, int val, int mode) {
	if (mode == 1) {
		for(int i = 0; i <= L; ++i) 
			res += bit[i][(((val >> i) & 1) ^ 1)] * 1LL * (1 << i);
		// cerr << "v: " << v << " res: " << res << '\n';
	}
	else {
		for(int i = 0; i <= L; ++i) 
			++bit[i][((val >> i) & 1)];
	}
    for (int u : g[v])
        if (u != p && !used[u])
            dfs(u, v, (val ^ w[u]), mode);
}

void calc(int v) {
	// cerr << "centr: " << v << '\n';
	calc_size(v, v);
	for(int i = 0; i <= L; ++i) 
		++bit[i][((w[v] >> i) & 1)];
	for(auto to : g[v]) {
		int val = w[to];
		if (!used[to]) {
			dfs(to, v, w[to], 1);
			dfs(to, v, (w[v]^w[to]), 2);
		}
	}
	for(int i = 0; i <= L; ++i) 
		bit[i][0] = bit[i][1] = 0;
	
	used[v] = 1;
	for(auto to : g[v]) {
		if(!used[to]) {
			calc(centroid(to, v, s[to]));
		}
	}
}

void solve() { 
	cin >> n;
	for(int i = 1; i <= n; ++i) {
		cin >> w[i];
		res += w[i];
	}
	for(int i = 1; i < n; ++i) {
		int u, v;
		cin >> u >> v;
		g[u].pb(v);
		g[v].pb(u);
	}
	calc_size(1, 1);
	calc(centroid(1, 1, n));
	cout << res << '\n';
}


signed main() {
	NFS;
	int test = 1;
//	cin >> test;
	for(int i = 1; i <= test; ++i){
		solve();
	}
}